import math 
import numpy as np
import matplotlib.pyplot as plt
import cv2

#图像灰度化
def toGrayImg(img):
    h,w = img.shape[:2]
    grayimg = np.zeros(img.shape[:2], dtype=np.uint8) 
    for i in range(h):
        for j in range(w):
            grayimg[i,j] = img[i,j,0] * 0.11 + img[i,j,1] * 0.59 + img[i,j,2] * 0.3
    return grayimg

#双线性插值放大缩小图片
def scrollImg(img, outimg):
    h,w,channel = img.shape
    print(f"srch:{h}, srcw:{w}")
    dsth,dstw = outimg[0], outimg[1]
    if h == dsth and w == dstw:
        return img.copy()
    dstimg = np.zeros([dsth,dstw,channel], dtype=np.uint8)
    scaley, scalex = float(h)/dsth, float(w)/dstw
    print(f"scalex:{scalex}, scaley:{scaley}")
    for i in range(channel):
        for dst_y in range(dsth):
            for dst_x in range(dstw):
                # print(f"x: {dst_x}, y:{dst_y}, channel:{i}")
                #中心对称后，求出目标图像在原始图像上的相对坐标
                srcx = (dst_x + 0.5) * scalex - 0.5
                srcy = (dst_y + 0.5) * scaley - 0.5
                #处理相对坐标，得到周围4个点相应的坐标位置
                src_x0 = int(np.floor(srcx))
                src_x1 = min(src_x0+1, w - 1)
                src_y0 = int(np.floor(srcy))
                src_y1 = min(src_y0+1, h - 1)
                #计算相应位置的像素值
                temp0 = (src_x1 - srcx) * img[src_y0, src_x0, i] + (srcx - src_x0) * img[src_y0, src_x1, i]
                temp1 = (src_x1 - srcx) * img[src_y1, src_x0, i] + (srcx - src_x0) * img[src_y1, src_x1, i]
                dstimg[dst_y, dst_x, i] = int((src_y1 - srcy) * temp0 + (srcy - src_y0) * temp1)
    return dstimg

#直方图均值化图片提亮增强，增加对比度
def promptImg(path):
    img = cv2.imread(path)
    cv2.imshow("src", img)

    # 彩色图像均衡化,需要分解通道 对每一个通道均衡化
    (b, g, r) = cv2.split(img)
    bH = cv2.equalizeHist(b)
    gH = cv2.equalizeHist(g)
    rH = cv2.equalizeHist(r)
    # 合并每一个通道
    result = cv2.merge((bH, gH, rH))
    cv2.imshow("dst_rgb", result)

    cv2.waitKey(0)

#直方图均值化，图像增强
def histogramEqualization(img):
    img = toGrayImg(img)
    newImg = np.zeros(img.shape[:2], dtype=np.uint8)
    h,w=img.shape[:2]
    sum = [0]*256
    out = [0]*256
    for i in range(h):
        for j in range(w):
            sum[img[i,j]] += 1
    sumpoint = 0
    for i in range(np.min(img),np.max(img) + 1):
        sumpoint += sum[i]
        out[i] = 256 * sumpoint/(h*w) - 1
    for i in range(h):
        for j in range(w):
            newImg[i,j] = int(round(out[img[i,j]]))
    # newImg = cv2.equalizeHist(img)
    return newImg

#高斯滤波，平滑图片
def gaussianFiltering(img, sigma, kernelSize):
    n1 = 1/(2*np.pi*sigma**2)
    n2 = -1/(2*sigma**2)
    h, w = img.shape[:2]
    gaussiankernel = np.zeros([kernelSize,kernelSize], dtype=np.float32)
    newImg = np.zeros(img.shape[:2], dtype=np.float32)
    #生成高斯核函数
    for i in range(kernelSize):
        for j in range(kernelSize):
            x = i - kernelSize//2
            y = j - kernelSize//2
            gaussiankernel[i,j] = n1*np.exp(n2*(x**2+y**2))
    gaussiankernel = gaussiankernel/np.sum(gaussiankernel) #归一化
    padsize = kernelSize//2
    padimg = np.pad(img, padsize, 'constant')
    #做卷积
    for i in range(h):
        for j in range(w):
            newImg[i,j] = np.sum(padimg[i:i+kernelSize,j:j+kernelSize]*gaussiankernel)
    cv2.imshow("gaussian img" ,newImg)
    return newImg

#sobel边缘检测，输出梯度图像
def sobelImg(img):
    h, w = img.shape[:2]
    sobelKernel_x = [[-1,0,1],
                     [-2,0,2],
                     [-1,0,1]]
    sobelKernel_y = [[1,2,1],
                     [0,0,0],
                     [-1,-2,-1]]
    dyImg = np.zeros(img.shape)
    dxImg = np.zeros(img.shape)
    tiduImg = np.zeros(img.shape)
    for i in range(1,h-1):
        for j in range(1,w-1):
            dyImg[i,j] = np.sum(sobelKernel_y * img[i-1:i+2,j-1:j+2])
            dxImg[i,j] = np.sum(sobelKernel_x * img[i-1:i+2,j-1:j+2])
            tiduImg[i,j] = np.sqrt(dxImg[i,j]**2  + dyImg[i,j]**2)
    dxImg[dxImg == 0] = 0.00000001
    angle = dyImg/dxImg

    # _, ret =  cv2.threshold(cv2.bitwise_or(dxImg,dyImg), 128, 255, cv2.THRESH_BINARY)
    return (tiduImg, angle)

#canny边缘检测算法(best)
def cannyEdgeImg(img, sigma, kernelsize):
    grayImg = toGrayImg(img)
    print(f'灰度化完成{grayImg.shape}')
    cv2.imshow('gray img', grayImg)
    h,w = grayImg.shape

    newImg = np.zeros(grayImg.shape)
    newImg = gaussianFiltering(grayImg, sigma, kernelsize)
    
    cv2.imshow('gaussian img',newImg.astype(np.uint8))
    print(f'高斯平滑完成{newImg.shape}')

    newImg, angle = sobelImg(grayImg)
    cv2.imshow('sobel img',newImg.astype(np.uint8))
    print(f'sobel边缘检测完成{newImg.shape}')

    dtemp = [0]*2
    try:
        for i in range(1,w-1):
            for j in range(1,h-1):
                # print(f"x:{i},y:{j}")
                if newImg[i,j] == 0:    #梯度为0的点直接归0
                    newImg[i,j] = 0
                else:
                    temp = newImg[i-1:i+2, j-1:j+2]  ##!!!左上角为(0,0)
                    if angle[i,j] >= 1:         #插值公式:f(x,y) = (x-x0)*f(x1,y1) + (x1-x)*f(x0,y0)
                        dtemp[0] = (temp[0,2] - temp[0,1])/angle[i,j] + temp[0,1]
                        dtemp[1] = (temp[2,0] - temp[2,1])/angle[i,j] + temp[2,1]
                    elif angle[i,j] > 0:
                        dtemp[0] = (temp[0,2] - temp[1,2])*angle[i,j] + temp[1,2]
                        dtemp[1] = (temp[2,0] - temp[1,0])*angle[i,j] + temp[1,0]
                    elif angle[i,j] <= -1:
                        dtemp[0] = (temp[0,1] - temp[0,0])/angle[i,j] + temp[0,1]
                        dtemp[1] = (temp[2,1] - temp[2,2])/angle[i,j] + temp[2,1]
                    elif angle[i,j] < 0:
                        dtemp[0] = (temp[1,0] - temp[0,0])*angle[i,j] + temp[1,0]
                        dtemp[1] = (temp[1,2] - temp[2,2])*angle[i,j] + temp[1,2]
                    if angle[i,j] != 0 and (newImg[i,j] < dtemp[0] or newImg[i,j] < dtemp[1]):
                        newImg[i,j] = 0
    except Exception as e:
        print(e)
        return newImg
    cv2.imshow('NMS img',newImg.astype(np.uint8))
    print(f'NMS非极大值抑制完成{newImg.shape}')

    low = newImg.mean()*0.5
    high = low * 3
    medium = []
    for i in range(1, w-1):
        for j in range(1, h-1):
            if newImg[i,j] < low:
                newImg[i,j] = 0
            elif newImg[i,j] > high:
                newImg[i,j] = 255
                medium.append([i,j])
    while not len(medium):
        x,y=medium.pop()
        a = newImg[x-1:x+2, y-1:y+2]
        if a[0,0] < high and a[0,0] > low:
            newImg[x-1,y-1] = 255
            medium.append([x-1, y-1])
        elif a[0,1] < high and a[0,1] > low:
            newImg[x-1,y] = 255
            medium.append([x-1, y])
        elif a[0,2] < high and a[0,2] > low:
            newImg[x-1,y+1] = 255
            medium.append([x-1, y+1])
        elif a[1,0] < high and a[1,0] > low:
            newImg[x,y-1] = 255
            medium.append([x, y-1])
        elif a[1,2] < high and a[1,2] > low:
            newImg[x,y+1] = 255
            medium.append([x, y+1])
        elif a[2,0] < high and a[2,0] > low:
            newImg[x+1,y-1] = 255
            medium.append([x+1, y-1])
        elif a[2,1] < high and a[2,1] > low:
            newImg[x+1,y] = 255
            medium.append([x+1, y])
        elif a[2,2] < high and a[2,2] > low:
            newImg[x+1,y+1] = 255
            medium.append([x+1, y+1])
    for i in range(newImg.shape[0]):
        for j in range(newImg.shape[1]):
            if newImg[i,j] != 0 and newImg[i,j] != 255:
                newImg[i,j] = 0
    print(f'阈值筛选完成{newImg.shape}')

    return newImg

#PCA主成分分析降维算法
def PCAImg(Img):
    return Img

if __name__ == '__main__':
    path = './images/sample2.jpg'
    oriImg = cv2.imread(path)
    print(f'原始图形{oriImg.shape}')
    cv2.imshow('ori', oriImg)

    newImg = cannyEdgeImg(oriImg, 0.5, 5)
    print(f"canny边缘检测完成{newImg.shape}")

    cv2.waitKey(0)